#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
import tqdm
import argparse
import onnx
from magiconnx import OnnxGraph


def find_data_flow(onnx_path, flow_list):

    flow_set = set(flow_list)
    graph = OnnxGraph(onnx_path)

    inits = graph.get_nodes("Initializer")
    init2value = {init.name: init.value for init in inits}
    init_names = set([init.name for init in inits])

    flow2node = {}
    op_names = set(graph._all_ops_name) - init_names
    for op_name in op_names:
        op = graph[op_name]
        if op._op_type == 'Constant' or op._op_type == 'Placeholder':
            continue
        for flow in op.inputs:
            if flow not in flow_set:
                continue
            if flow in flow2node:
                flow2node[flow][1].append(op_name)
            else:
                flow2node[flow] = [[], [op_name]]
        for flow in op.outputs:
            if flow not in flow_set:
                continue
            if flow in flow2node:
                flow2node[flow][0].append(op_name)
            else:
                flow2node[flow] = [[op_name], []]

    for flow in flow_list:
        if flow in flow2node:
            input_nodes = flow2node[flow][0]
            output_nodes = flow2node[flow][1]
            print(f"{flow}  {input_nodes} --> {output_nodes}")
        elif flow in init2value:
            value = str(init2value[flow]).replace('\n', ' ')
            if len(value) > 20:
                print(f"{flow}  Constant")
            else:
                print(f"{flow}  Constant({value})")
        else:
            print(f"{flow}  Uknown")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--onnx-path', type=str, required=True, help='.')
    parser.add_argument('--flow-names', type=str, nargs='+', required=True, help='list of data flow names')
    args = parser.parse_args()

    find_data_flow(args.onnx_path, args.flow_names)
